import copy
from pkl import *
from strategy import *
strategy=strategy()
pkl=pkl()
stock=get_stock()
# 交易回测类
class trade:
    def __init__(self,startdate,enddate,asset,data):
        self.name='交易方法与回测过程'
        self.stock = {}
        start = startdate
        end = enddate
        for code in data.keys():
            for date in data[code].index.tolist():
                if not np.isnan(data[code].loc[date, 'close']):
                    closepricelist = data[code].loc[date:"2019-12-31"]
                    break
            s = strategy.mean_reversion(closepricelist)
            list1 = s
            buyingpoint = list1[0]
            sellingpoint = list1[1]
            n = 24
            total = float(asset) /n
            t = self.Trade(start, end, buyingpoint, sellingpoint, data[code], total)
            self.stock[code] = copy.deepcopy(t)


    # 按照移动均线策略分别对所选的全部股票进行交易
    def trade_all(self,trade_dict,capital,start_date,end_date):
        st={}
        # 股票数量
        n=len(trade_dict)
        # 每只股票投入的资金
        money=float(capital)/n
        for code in trade_dict.keys():
            data=stock.get_a_stock(code,start_date,end_date)
            buy_time=trade_dict[code][0]
            sell_time=trade_dict[code][1]
            t = self.Trade(start_date,end_date, buy_time, sell_time, data, money)
            st[code] = copy.deepcopy(t)
        return st

    def Trade(self, starttime, endtime, buyingpoint, sellingpoint, stockinfo, total):
        account = pd.DataFrame(columns=['stock', 'cash', 'shares', 'total'])
        cash = float(total)
        stockvalue = 0
        shares = 0
        stockinfo = stockinfo.fillna(0)
        from datetime import datetime  # 看看能不能去掉
        year, mon, day = starttime.split('-')
        start = datetime(int(year), int(mon), int(day))
        year, mon, day = endtime.split('-')
        end = datetime(int(year), int(mon), int(day))

        for date in stockinfo.index.tolist():
            year, mon, day = date.split('-')
            today = datetime(int(year), int(mon), int(day))
            if today >= start and today <= end:
                openprice = float(stockinfo.loc[date, 'open'])
                closeprice = float(stockinfo.loc[date, 'close'])
                rate = pow(1.036, 1 / 365) - 1
                if date in list(buyingpoint):
                    if cash >= openprice * 100 * (1 + 0.002):
                        shares = shares + (cash // (openprice * 100 * (1 + 0.002))) * 100
                        stockvalue = shares * closeprice
                        cash = (cash % (openprice * 100 * (1 + 0.002))) * (1 + rate)
                    else:
                        stockvalue = shares * closeprice
                        cash = cash * (1 + rate)
                if date in list(sellingpoint):
                    cash = (cash + shares * openprice) * (1 + rate)
                    shares = 0
                    stockvalue = 0
                if date not in list(buyingpoint) and date not in list(sellingpoint):
                    stockvalue = shares * closeprice
                    cash = cash * (1 + rate)
                account.loc[date, 'cash'] = cash
                account.loc[date, 'stock'] = stockvalue
                account.loc[date, 'shares'] = shares
                account.loc[date, 'total'] = cash + stockvalue
        print(account)

        return account

    # 前两个交易策略的交易数据处理方法
class trader:

    def __init__(self,date_list, stock_dict, data, asset, period, num):
        asset = float(asset)
        period = int(period)
        num = int(num)
        # 按股票个数均分初始资金
        share = asset / num
        day = date_list[0]
        stocks = stock_dict[day]
        total_list = []
        rate_list = [0]
        for i in range(len(date_list)):
            total = 0
            date = date_list[i]

            for s in stocks:
                price = float(data[date].close[s])
                price0 = float(data[day].close[s])
                stock_share = share * (price / price0)
                total += stock_share

            if i % period == 0:
                share = total / num
                day = date
                stocks = stock_dict[date]

            total_list.append(total)
            if i != 0:
                rate = (total_list[-1] - total_list[-2]) / total_list[-2]
                rate_list.append(rate)

        account_df = pd.DataFrame({"总资本": total_list,
                                   "日收益率": rate_list
                                   })

        account_df.index = ["%s" % x for x in date_list]
        print(account_df)
        total_rate=(total-asset)/asset
        print('总收益率：',total_rate)
        account_df.to_csv('./strategy.csv')
        print('回测数据已保存到strategy.csv')











